#!/usr/bin/env python3
"""
Neural Trajectory Analyzer for Embodied World Models
分析稳定周期中的神经轨迹
"""

from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import matplotlib.pyplot as plt
import time
from datetime import datetime
import numpy as np
from sklearn.decomposition import PCA
from mpl_toolkits.mplot3d import Axes3D
import sys

@partial(jax.jit, static_argnums=(3,))
def model_forward(variables, state, x, model):
    """ forward pass of the model """
    return model.apply(variables, state, x)

@jit
def get_action(y):
    return jnp.argmax(y)
get_action_vmap = jax.vmap(get_action)

def load_task(pth = "./logs/task.json"):
    # open json file
    with open(pth, "r") as f:
        data = json.load(f)
        landscape = data["data"]
        state = data["state"]
        goal = data["goal"]
        print("Task loaded - Start:", state, "Goal:", goal)
    return landscape, state, goal

def progress_bar(current, total, width=50):
    """显示进度条"""
    percent = float(current) * 100 / total
    filled = int(width * current // total)
    bar = '█' * filled + '-' * (width - filled)
    sys.stdout.write(f'\r[{bar}] {percent:.1f}% ({current}/{total})')
    sys.stdout.flush()

def segment_neural_trajectories(neural_states, goal_reached_history):
    """
    根据目标到达历史分割神经轨迹为完整周期
    
    Args:
        neural_states: 完整的神经状态序列 [time_steps, neural_dim]
        goal_reached_history: 目标到达历史 [time_steps]
    
    Returns:
        cycles: 周期列表，每个周期包含从起点到终点的神经状态
    """
    cycles = []
    cycle_start = 0
    
    for i, goal_reached in enumerate(goal_reached_history):
        if goal_reached:
            # 找到一个完整周期：从cycle_start到i
            if i > cycle_start:
                cycle_neural_states = neural_states[cycle_start:i+1]
                cycles.append(cycle_neural_states)
                print(f"Cycle {len(cycles)}: steps {cycle_start}-{i} (length: {len(cycle_neural_states)})")
            cycle_start = i + 1
    
    return cycles

def analyze_neural_trajectories(cycles, num_cycles=10):
    """
    分析最后一个周期的神经轨迹
    
    Args:
        cycles: 所有周期的神经状态列表
        num_cycles: 分析的周期数量（此版本只使用最后一个）
    
    Returns:
        pca_result: PCA分析结果
        cycle_data: 选取的周期数据
    """
    if len(cycles) == 0:
        print("Error: No cycles found!")
        selected_cycles = []
    else:
        selected_cycles = [cycles[-1]]  # 只选择最后一个周期
        print(f"Analyzing last cycle (length: {len(cycles[-1])})")
    
    # 将最后一个周期的神经状态准备分析
    all_neural_states = []
    cycle_labels = []
    
    for cycle_idx, cycle in enumerate(selected_cycles):
        all_neural_states.extend(cycle)
        cycle_labels.extend([cycle_idx] * len(cycle))
    
    all_neural_states = np.array(all_neural_states)
    cycle_labels = np.array(cycle_labels)
    
    print(f"Total neural states for analysis: {all_neural_states.shape}")
    
    # 执行PCA分析
    pca = PCA(n_components=3)
    pca_result = pca.fit_transform(all_neural_states)
    
    print(f"PCA explained variance ratio: {pca.explained_variance_ratio_}")
    print(f"Cumulative explained variance: {np.cumsum(pca.explained_variance_ratio_)}")
    
    return pca_result, cycle_labels, selected_cycles, pca

def visualize_neural_trajectories_2d(pca_result, cycle_labels, selected_cycles, pca, perturbation_info=None):
    """
    2D可视化神经轨迹 - 显示闭合的极限环和扰动向量（前两个主成分）
    """
    fig, ax = plt.subplots(figsize=(12, 10))
    
    # 为每个周期使用不同颜色
    colors = plt.cm.tab10(np.linspace(0, 1, len(selected_cycles)))
    
    # 绘制每个周期的轨迹
    current_idx = 0
    for cycle_idx, cycle in enumerate(selected_cycles):
        cycle_length = len(cycle)
        cycle_pca = pca_result[current_idx:current_idx + cycle_length]
        
        # 创建闭合轨迹：将最后一点连接到第一点
        closed_cycle_pca = np.vstack([cycle_pca, cycle_pca[0:1]])  # 添加起点到末尾形成闭合
        
        # 绘制状态转移箭头，显示从一个状态到下一个状态的方向
        for i in range(len(closed_cycle_pca) - 1):
            start_point = closed_cycle_pca[i]
            end_point = closed_cycle_pca[i + 1]
            
            # 计算箭头方向
            dx = end_point[0] - start_point[0]
            dy = end_point[1] - start_point[1]
            
            # 绘制状态转移箭头
            ax.arrow(start_point[0], start_point[1], dx, dy,
                    head_width=0.15, head_length=0.15, 
                    fc=colors[cycle_idx], ec=colors[cycle_idx],
                    linewidth=2, alpha=0.8, length_includes_head=True)
        
        # 为最后一个周期添加label
        if cycle_idx == 0:
            ax.plot([], [], color=colors[cycle_idx], linewidth=2, alpha=0.8, 
                   label='State Transitions (Last Cycle)')
        
        # 绘制轨迹点，显示时间演进
        for i in range(cycle_length):
            alpha = 0.3 + 0.7 * (i / cycle_length)  # 透明度渐变显示时间演进
            ax.scatter(cycle_pca[i, 0], cycle_pca[i, 1], 
                      color=colors[cycle_idx], s=50, alpha=alpha)
        
        # 特别标记起点（大圆）和终点（方块）
        ax.scatter(cycle_pca[0, 0], cycle_pca[0, 1], 
                  color=colors[cycle_idx], s=200, marker='o', alpha=1.0, 
                  edgecolor='black', linewidth=2, label='Start')
        ax.scatter(cycle_pca[-1, 0], cycle_pca[-1, 1], 
                  color=colors[cycle_idx], s=200, marker='s', alpha=1.0,
                  edgecolor='black', linewidth=2, label='End')
        
        current_idx += cycle_length
    
    ax.set_xlabel('PC1 (Explained Variance: {:.1f}%)'.format(pca.explained_variance_ratio_[0]*100))
    ax.set_ylabel('PC2 (Explained Variance: {:.1f}%)'.format(pca.explained_variance_ratio_[1]*100))
    ax.set_title('Neural Limit Cycles in Stable Phase\n(Closed trajectories in PC1-PC2 space)')
    ax.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    ax.grid(True, alpha=0.3)
    
    # 绘制扰动向量（已注释掉）
    # if perturbation_info is not None:
    #     perturbation_step, perturbation_vector, perturbation_state = perturbation_info
    #     
    #     # 将扰动前的神经状态转换到PCA空间
    #     perturbation_state_pca = pca.transform(perturbation_state.reshape(1, -1))[0]
    #     
    #     # 将扰动向量转换到PCA空间
    #     perturbation_vector_pca = pca.transform((perturbation_state + perturbation_vector).reshape(1, -1))[0]
    #     perturbation_direction_pca = perturbation_vector_pca - perturbation_state_pca
    #     
    #     # 绘制扰动位置（红色大点）
    #     ax.scatter(perturbation_state_pca[0], perturbation_state_pca[1], 
    #               color='red', s=300, marker='X', alpha=1.0, 
    #               edgecolor='darkred', linewidth=3, 
    #               label=f'Perturbation Point (Step {perturbation_step})')
    #     
    #     # 绘制扰动向量（红色箭头）
    #     ax.arrow(perturbation_state_pca[0], perturbation_state_pca[1],
    #             perturbation_direction_pca[0], perturbation_direction_pca[1],
    #             head_width=0.1, head_length=0.1, fc='red', ec='darkred',
    #             linewidth=3, alpha=0.8, label='Perturbation Vector')
    #     
    #     # 添加扰动信息文本
    #     ax.text(0.02, 0.98, f'Perturbation at step {perturbation_step}\nMagnitude: {np.linalg.norm(perturbation_vector):.4f}', 
    #             transform=ax.transAxes, verticalalignment='top',
    #             bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8))
    
    # 设置相等的坐标轴比例
    ax.set_aspect('equal', adjustable='box')
    
    plt.tight_layout()
    
    # 根据是否有扰动选择保存文件名
    filename = './logs/neural_trajectories_with_perturbation_2d.png' if perturbation_info else './logs/neural_trajectories_2d.png'
    plt.savefig(filename, dpi=300, bbox_inches='tight')
    plt.show()

def main():
    """ Neural Trajectory Analysis Main Function """
    print("=" * 60)
    print("   Neural Trajectory Analyzer")
    print("   分析具身世界模型中的稳定周期神经轨迹")
    print("=" * 60)
    
    # Parse arguments
    rpl_config = ReplayConfig()
    parser = argparse.ArgumentParser(description='Analyze neural trajectories in stable cycles')
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--life_duration", type=int, default=1000)
    parser.add_argument("--num_cycles", type=int, default=10, help="Number of cycles to analyze")

    args = parser.parse_args()

    # Load model
    print("Loading model...")
    params = load_weights(args.model_pth)

    # Create agent
    if args.nn_type == "vanilla":
        model = RNN(hidden_dims=args.nn_size)
    elif args.nn_type == "gru":
        model = GRU(hidden_dims=args.nn_size)

    # Load task
    if os.path.isfile(args.task_pth):
        landscape, state, goal = load_task(pth=args.task_pth)
        landscape = [landscape]
    else:
        print("Error: Task file not found")
        return

    # Create environment
    print("Creating environment...")
    GE = GridEnv(landscapes=landscape, width=12, height=12, num_envs_per_landscape=1)
    GE.reset()
    
    # Set initial states
    GE.batched_states = GE.batched_states.at[0, 0].set(state[0])
    GE.batched_states = GE.batched_states.at[0, 1].set(state[1])
    GE.batched_goals = GE.batched_goals.at[0, 0].set(goal[0])
    GE.batched_goals = GE.batched_goals.at[0, 1].set(goal[1])
    GE.init_batched_states, GE.init_batched_goals = jnp.copy(GE.batched_states), jnp.copy(GE.batched_goals)
    GE.batched_goal_reached = batch_compute_goal_reached(GE.batched_states, GE.batched_goals)
    GE.last_batched_goal_reached = jnp.copy(GE.batched_goal_reached)
    GE.concat_obs = get_ideal_obs_vmap(GE.batched_envs, GE.batched_states, GE.batched_goals, GE.last_batched_goal_reached)

    concat_obs = GE.concat_obs
    rnn_state = model.initial_state(GE.num_envs)

    # Data collection with perturbation
    print("Starting neural state collection with perturbation...")
    neural_states_history = []
    goal_reached_history = []
    render_id = 0
    
    # 扰动相关变量
    perturbation_N = 100  # 最后N步中施加扰动
    perturbation_step = args.life_duration - np.random.randint(1, perturbation_N + 1)  # 随机选择扰动位置
    perturbation_vector = None
    perturbation_state_before = None
    perturbation_applied = False
    
    print(f"Perturbation will be applied at step {perturbation_step}")

    # Run simulation
    for step in range(args.life_duration):
        # Progress display
        if step % 10 == 0:
            progress_bar(step, args.life_duration)

        # Model forward pass
        rnn_state, y1 = model_forward(params, rnn_state, concat_obs, model)
        
        # 施加扰动
        if step == perturbation_step and not perturbation_applied:
            # 记录扰动前的神经状态
            perturbation_state_before = np.array(rnn_state[0])
            
            # 生成扰动向量（-0.1到0.1的均匀分布）
            perturbation_vector = np.random.uniform(-0.3, 0.3, size=rnn_state[0].shape)
            
            # 施加扰动到神经状态
            # rnn_state = rnn_state.at[0].add(perturbation_vector)
            perturbation_applied = True
            
            print(f"\nPerturbation applied at step {step}")
            print(f"Perturbation magnitude: {np.linalg.norm(perturbation_vector):.4f}")

        batched_actions = get_action_vmap(y1)

        # Environment step
        batched_goal_reached, concat_obs = GE.step(batched_actions, reset=True)

        # Record neural state and goal reached status
        neural_states_history.append(np.array(rnn_state[0]))
        goal_reached_history.append(bool(batched_goal_reached[render_id]))

    print("\nData collection completed!")
    
    # Convert to numpy arrays
    neural_states_history = np.array(neural_states_history)
    goal_reached_history = np.array(goal_reached_history)
    
    print(f"Collected {len(neural_states_history)} neural states")
    print(f"Goal reached {np.sum(goal_reached_history)} times")

    # Segment into cycles
    print("\nSegmenting neural trajectories into cycles...")
    cycles = segment_neural_trajectories(neural_states_history, goal_reached_history)
    
    if len(cycles) == 0:
        print("Error: No complete cycles found!")
        return
    
    print(f"Found {len(cycles)} complete cycles")
    
    # Analyze trajectories
    print("\nAnalyzing neural trajectories...")
    pca_result, cycle_labels, selected_cycles, pca = analyze_neural_trajectories(cycles, args.num_cycles)
    
    # 准备扰动信息
    perturbation_info = None
    if perturbation_applied and perturbation_vector is not None:
        perturbation_info = (perturbation_step, perturbation_vector, perturbation_state_before)
    
    # Visualize
    print("\nGenerating 2D visualization with perturbation...")
    visualize_neural_trajectories_2d(pca_result, cycle_labels, selected_cycles, pca, perturbation_info)
    
    filename = "neural_trajectories_with_perturbation_2d.png" if perturbation_info else "neural_trajectories_2d.png"
    print(f"\nAnalysis completed! Results saved to ./logs/{filename}")

if __name__ == "__main__":
    main()
